// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/vector/base_vector.h>

namespace kumo::pollux::dwio::common {
    template<typename T>
    inline void ensureCapacity(
        BufferPtr &data,
        size_t numElements,
        pollux::memory::MemoryPool *pool,
        bool preserveOldData = false,
        bool clearBits = false) {
        size_t oldSize = 0;
        size_t newCapacity = BaseVector::byteSize<T>(numElements);
        if (!data) {
            data = AlignedBuffer::allocate<T>(numElements, pool);
        } else {
            oldSize = data->size();
            if (!data->isMutable() || data->capacity() < newCapacity) {
                auto newData = AlignedBuffer::allocate<T>(numElements, pool);
                if (preserveOldData) {
                    std::memcpy(
                        newData->template asMutable<uint8_t>(),
                        data->as<uint8_t>(),
                        oldSize);
                }
                data = newData;
            }
        }

        if (clearBits && newCapacity > oldSize) {
            std::memset(
                (void *) (data->asMutable<int8_t>() + oldSize),
                0L,
                newCapacity - oldSize);
        }
    }

    template<typename T>
    inline T *resetIfWrongVectorType(VectorPtr &result) {
        if (result) {
            auto casted = result->as<T>();
            // We only expect vector to be used by a single thread.
            if (casted && result.use_count() == 1) {
                return casted;
            }
            result.reset();
        }
        return nullptr;
    }

    template<typename... T>
    inline void resetIfNotWritable(VectorPtr &result, T &... buffer) {
        // The result vector and the buffer both hold reference, so refCount is at
        // least 2
        auto resetIfShared = [](auto &buffer) {
            const bool reset = buffer->refCount() > 2;
            if (reset) {
                buffer.reset();
            }
            return reset;
        };

        if ((... | resetIfShared(buffer))) {
            result.reset();
        }
    }
} // namespace kumo::pollux::dwio::common
